package org.jeecg;

import org.geotools.data.DataStore;
import org.geotools.data.DataStoreFinder;
import org.geotools.data.simple.SimpleFeatureCollection;
import org.geotools.data.simple.SimpleFeatureIterator;
import org.locationtech.jts.geom.MultiPolygon;
import org.locationtech.jts.geom.Polygon;
import org.locationtech.jts.io.WKBWriter;
import org.opengis.feature.simple.SimpleFeature;
import org.opengis.feature.simple.SimpleFeatureType;


import java.io.File;
import java.sql.*;
import java.util.Map;
import java.util.UUID;

public class Shp2PostGIS {

    public static void main(String[] args) {
        String shpPath = "F:\\地质数据\\地形坡度\\地形坡度底图.shp";

        // PostgreSQL JDBC配置
        String jdbcUrl = "jdbc:postgresql://123.56.202.52:5432/gisdb";
        String username = "postgres";
        String password = "wwq20010215";
        String tableName = "geotype";

        Connection conn = null;
        PreparedStatement pstmt = null;
        SimpleFeatureIterator iterator = null;
        WKBWriter wkbWriter = new WKBWriter(); // WKB写入器，复用提高效率

        try {
            long programStart = System.currentTimeMillis();

            // 1. 加载Shapefile
            File shpFile = new File(shpPath);
            DataStore shpDataStore = DataStoreFinder.getDataStore(Map.of(
                    "url", shpFile.toURI().toURL(),
                    "create spatial index", Boolean.FALSE
            ));
            if (shpDataStore == null) throw new RuntimeException("无法加载Shapefile数据源");

            String typeName = shpDataStore.getTypeNames()[0];
            SimpleFeatureCollection features = shpDataStore.getFeatureSource(typeName).getFeatures();
            SimpleFeatureType sourceType = shpDataStore.getSchema(typeName);
            // 2. 打印结构用于调试
            System.out.println("Shapefile字段信息:");
            sourceType.getAttributeDescriptors().forEach(attr ->
                    System.out.println("  " + attr.getLocalName() + " : " +
                            attr.getType().getBinding().getSimpleName()));
            long featureCount = features.size();
            System.out.println("找到 " + featureCount + " 条记录");

            // 2. 连接PostgreSQL
            conn = DriverManager.getConnection(jdbcUrl, username, password);
            conn.setAutoCommit(false); // 禁用自动提交

            // 3. 创建表（简化版，只包含需要的字段）
            createTable(conn, tableName);
            System.out.println("表 " + tableName + " 创建成功");

            // 4. 准备批量插入语句
            String insertSQL = "INSERT INTO " + tableName + " (id, the_geom, id_orig, gridcode) " +
                    "VALUES (?, ST_GeomFromWKB(?, 32650), ?, ?)";
            pstmt = conn.prepareStatement(insertSQL);

            // 5. 导入数据
            System.out.println("开始导入数据...");
            iterator = features.features();
            final int BATCH_SIZE = 10000;  // 批量大小
            int totalCount = 0;
            int currentBatch = 0;
            long batchStartTime = System.currentTimeMillis();
            long lastPrintTime = System.currentTimeMillis();

            while (iterator.hasNext()) {
                SimpleFeature feature = iterator.next();

                // 设置参数
                pstmt.setObject(1, UUID.randomUUID()); // id (主键)

                // 几何对象 - 直接处理为MultiPolygon
                MultiPolygon geom = (MultiPolygon) feature.getAttribute("the_geom");
                pstmt.setBytes(2, wkbWriter.write(geom)); // WKB格式

                // Id (Long)
                Long idOrig = (Long) feature.getAttribute("Id");
                pstmt.setLong(3, idOrig != null ? idOrig : -1); // 处理空值

                // gridcode (Long)
                Long gridcode = (Long) feature.getAttribute("gridcode");
                pstmt.setLong(4, gridcode != null ? gridcode : -1);

                pstmt.addBatch();
                totalCount++;
                currentBatch++;

                // 定期提交批次
                if (currentBatch >= BATCH_SIZE) {
                    executeBatch(conn, pstmt, currentBatch, batchStartTime);
                    currentBatch = 0;
                    batchStartTime = System.currentTimeMillis();
                }

                // 每5秒打印一次进度
                long currentTime = System.currentTimeMillis();
                if (currentTime - lastPrintTime > 5000) {
                    System.out.printf("总进度: %d/%d 条 (%.1f%%)%n",
                            totalCount, featureCount,
                            (totalCount * 100.0) / featureCount);
                    lastPrintTime = currentTime;
                }
            }

            // 执行最后一批
            if (currentBatch > 0) {
                executeBatch(conn, pstmt, currentBatch, batchStartTime);
            }

            System.out.println("数据导入完成!");

            // 输出总耗时
            long importDuration = System.currentTimeMillis() - programStart;
            System.out.println("========================================");
            System.out.printf("导入完成! 共导入 %d 条记录, 总耗时 %d ms (%.1f 分钟)%n",
                    totalCount, importDuration, importDuration / 60000.0);
            System.out.printf("平均速率: %.0f 条/秒%n", totalCount * 1000.0 / importDuration);
            System.out.println("========================================");

            // 6. 打印索引创建SQL
            System.out.println("\n请执行以下SQL创建空间索引:");
            System.out.println("-----------------------------------------------------");
            System.out.println(getSpatialIndexSQL(tableName, "the_geom"));
            System.out.println("-----------------------------------------------------");
            System.out.println("注意: 创建索引可能需要较长时间，建议在数据库空闲时执行");

            // 最终统计
            long totalProgramTime = System.currentTimeMillis() - programStart;
            System.out.println("========================================");
            System.out.println("程序执行完成!");
            System.out.printf("总耗时: %d ms (%.1f 分钟)%n",
                    totalProgramTime, totalProgramTime / 60000.0);
            System.out.println("========================================");

        } catch (Exception e) {
            e.printStackTrace();
            System.err.println("导入失败: " + e.getMessage());
            try {
                if (conn != null) conn.rollback();
            } catch (SQLException ex) {
                System.err.println("回滚失败: " + ex.getMessage());
            }
        } finally {
            // 关闭所有资源
            try {
                if (iterator != null) iterator.close();
                if (pstmt != null) pstmt.close();
                if (conn != null) conn.close();
            } catch (Exception e) {
                e.printStackTrace();
            }
        }
    }

    // 创建表（简化版）
    private static void createTable(Connection conn, String tableName)
            throws SQLException {
        try (Statement stmt = conn.createStatement()) {
            // 删除现有表（如果存在）
            stmt.execute("DROP TABLE IF EXISTS " + tableName + " CASCADE");

            // 创建新表，只包含需要的字段
            String createSQL = "CREATE TABLE " + tableName + " (" +
                    "id UUID PRIMARY KEY, " +
                    "the_geom GEOMETRY(MultiPolygon, 32650), " +
                    "id_orig BIGINT, " + // 原始Id字段
                    "gridcode BIGINT)";  // gridcode字段

            stmt.execute(createSQL);
            conn.commit();
        }
    }

    // 执行批处理
    private static void executeBatch(Connection conn, PreparedStatement pstmt, int batchSize, long batchStartTime)
            throws SQLException {
        long start = System.currentTimeMillis();
        pstmt.executeBatch();
        conn.commit();
        long duration = System.currentTimeMillis() - batchStartTime;

        // 计算当前速率（条/秒）
        double rate = batchSize * 1000.0 / duration;

        System.out.printf("已提交批次: %d条, 耗时: %d ms, 速率: %.0f 条/秒%n",
                batchSize, duration, rate);
    }

    // 获取空间索引创建SQL
    private static String getSpatialIndexSQL(String table, String geomColumn) {
        return "CREATE INDEX CONCURRENTLY " + table + "_" + geomColumn + "_spatial_idx\n" +
                "ON " + table + "\n" +
                "USING GIST (" + geomColumn + ");\n\n" +
                "-- 可选: 分析表以优化查询性能\n" +
                "ANALYZE " + table + ";";
    }
}